package org.base23.ai.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;

@Configuration
@Slf4j
public class VectorStoreConfig {

    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return new SimpleVectorStore(embeddingModel);
    }

    @Bean
    CommandLineRunner ingestTermOfServiceToVectorStore(
            EmbeddingModel embeddingModel,
            VectorStore vectorStore,
            @Value("classpath:rag/*.txt") Resource[] termsOfServiceDocs) {
        return args -> {
            try {
                for (Resource doc : termsOfServiceDocs) {
                    if (doc.exists()) {
                        log.info("RAG正在加载文件: {}", doc.getFilename());
                        vectorStore.write(
                                new TokenTextSplitter().transform(
                                        new TextReader(doc).read())
                        );
                        log.info("成功加载: {}", doc.getFilename());
                    } else {
                        log.warn("文件未找到: {}", doc.getFilename());
                    }
                }
            } catch (Exception e) {
                log.error("向量数据库写入失败: {}", e.getMessage());
            }
        };
    }
}
